import cv2
import numpy as np
from sklearn.cluster import KMeans


class ClothingColorExtractor:

    def analyze_image(self, img_path):
        img = cv2.imread(img_path)
        if img is None:
            raise Exception(f"无法读取图片: {img_path}")

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        masked_img = self.clothing_detection_sam(img)
        colors, percentages = self.extract_dominant_colors(masked_img, self.color_num_var.get())

        return {
            'original_img': img,
            'masked_img': masked_img,
            'colors': colors,
            'percentages': percentages
        }

    def clothing_detection_sam(self, img):
        self.predictor.set_image(img)
        h, w = img.shape[:2]
        input_point = np.array([[w // 2, h // 2]])
        input_label = np.array([1])
        masks, scores, logits = self.predictor.predict(
            point_coords=input_point,
            point_labels=input_label,
            multimask_output=True,
        )
        mask = masks[np.argmax(scores)]
        masked_img = img * mask[:, :, np.newaxis]
        return masked_img

    def extract_dominant_colors(self, img, num_colors=3):
        pixels = img.reshape(-1, 3)
        pixels = pixels[~np.all(pixels == 0, axis=1)]

        if len(pixels) == 0:
            return [(0, 0, 0)], [1.0]

        kmeans = KMeans(n_clusters=num_colors, random_state=42)
        kmeans.fit(pixels)

        colors = kmeans.cluster_centers_.astype(int)
        labels = kmeans.labels_
        percentages = np.bincount(labels) / len(labels)

        sorted_indices = np.argsort(percentages)[::-1]
        colors = colors[sorted_indices]
        percentages = percentages[sorted_indices]

        return colors, percentages

    def create_color_mask(self, img, target_color, threshold=50):
        img = img.astype(int)
        target_color = np.array(target_color, dtype=int)
        diff = np.abs(img - target_color)
        mask = np.all(diff <= threshold, axis=2)
        return mask